import os
import pandas as pd
import numpy as np
import torch
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
from PIL import Image
from sklearn.preprocessing import LabelEncoder, OneHotEncoder

class Derm7PtDataset(Dataset):
    """ PyTorch Dataset for Derm7Pt """

    def __init__(self, df, transform=None):
        """
        初始化数据集
        :param df: 预处理后的 Pandas DataFrame
        :param transform: 图像预处理
        """
        self.df = df
        self.transform = transform

    def __len__(self):
        """ 返回数据集大小 """
        return len(self.df)

    def __getitem__(self, idx):
        """ 取单个样本 """
        row = self.df.iloc[idx]

        # 加载临床图像
        clinic_img = Image.open(row["clinic_path"]).convert("RGB")
        derm_img = Image.open(row["derm_path"]).convert("RGB")

        # 进行变换
        if self.transform:
            clinic_img = self.transform(clinic_img)
            derm_img = self.transform(derm_img)

        # 加载数值特征
        metadata = torch.tensor(row["metadata"], dtype=torch.float32)

        # 目标类别（分类标签）
        label = torch.tensor(row["diagnosis_encoded"], dtype=torch.long)

        return clinic_img, derm_img, metadata, label

def get_dataloader(data_dir, batch_size=32):
    """
    预处理 Derm7Pt 数据集并返回 PyTorch DataLoader
    :param data_dir: 数据目录
    :param batch_size: batch size
    :return: train_loader, valid_loader, test_loader
    """
    # 加载 meta 数据
    meta_path = os.path.join(data_dir, "meta/meta.csv")
    df = pd.read_csv(meta_path)

    # 处理图像路径
    df["clinic_path"] = df["clinic"].apply(lambda x: os.path.join(data_dir, "images", x))
    df["derm_path"] = df["derm"].apply(lambda x: os.path.join(data_dir, "images", x))

    # 处理分类标签
    label_encoder = LabelEncoder()
    df["diagnosis_encoded"] = label_encoder.fit_transform(df["diagnosis"])

    # 处理 7 点检查法评分（标准化）
    seven_point_cols = [
        "pigment_network", "streaks", "pigmentation", "regression_structures",
        "dots_and_globules", "blue_whitish_veil", "vascular_structures"
    ]
    for col in seven_point_cols:
        df[col] = df[col].map({"absent": 0, "present": 1}).fillna(0)

    # 处理元数据（性别、位置、病变形态）
    categorical_cols = ["sex", "location", "elevation", "management"]
    onehot_encoder = OneHotEncoder(sparse=False, handle_unknown="ignore")
    encoded_metadata = onehot_encoder.fit_transform(df[categorical_cols])

    # 组合元数据
    df["metadata"] = list(encoded_metadata)

    # 读取数据集划分索引
    train_idx = pd.read_csv(os.path.join(data_dir, "meta/train_indexes.csv"))["indexes"].tolist()
    valid_idx = pd.read_csv(os.path.join(data_dir, "meta/valid_indexes.csv"))["indexes"].tolist()
    test_idx = pd.read_csv(os.path.join(data_dir, "meta/test_indexes.csv"))["indexes"].tolist()

    # 划分数据集
    df_train = df.iloc[train_idx].reset_index(drop=True)
    df_valid = df.iloc[valid_idx].reset_index(drop=True)
    df_test = df.iloc[test_idx].reset_index(drop=True)

    # 图像转换
    image_transform = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
    ])

    # 创建 Dataset
    train_dataset = Derm7PtDataset(df_train, transform=image_transform)
    valid_dataset = Derm7PtDataset(df_valid, transform=image_transform)
    test_dataset = Derm7PtDataset(df_test, transform=image_transform)

    # 创建 DataLoader
    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True, num_workers=4)
    valid_loader = DataLoader(valid_dataset, batch_size=batch_size, shuffle=False, num_workers=4)
    test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False, num_workers=4)

    return train_loader, valid_loader, test_loader
